package cn.yuemouren.tree.trie;

import java.util.TreeMap;

/**
 * @Author: Timi
 * @Description: tri(多叉树)
 *     扩展：压缩树,Term tree,后缀树, 包含的树算法
 * @Date: 2020/7/11 15:15
 * @Version: v1.0
 */
public class Trie {

    private class Node {

        private boolean isLeaf;//是否是节点
        private TreeMap<Character,Node> next;//目前节点对应下面的节点

        public Node() {
            this(false);
        }

        public Node(boolean isLeaf) {
            this.isLeaf = isLeaf;
            this.next = new TreeMap<>();
        }
    }

    private Node root;//根节点
    private int size;//存在的单词数量

    public Trie() {
        this.root = new Node();
        this.size = 0;
    }

    /**
     * 向trie中插入数据
     * 非递归写法
     * @param word
     */
    public void add(String word){
        Node cur = root;
        for (int i = 0; i < word.length(); i++){
            Character character = word.charAt(i);
            if(null == cur.next.get(character)){
                cur.next.put(character,new Node());
            }
            cur = cur.next.get(character);
        }
        if(!cur.isLeaf){
            cur.isLeaf = true;
            size++;
        }
    }

    /**
     * 向trie中插入数据
     * 递归写法
     * @param word
     */
    public void addRecursive(String word){
        addRecursive(word,0,root);
    }

    /**
     * 根据index 获取到字符，然后判断节点中是否包含
     * @param word
     * @param index
     * @param cur
     */
    private void addRecursive(String word,int index,Node cur){
        if (index >= word.length()){
            if(!cur.isLeaf){
                cur.isLeaf = true;
                size++;
            }
            return;
        }
        if(null == cur.next.get(word.charAt(index))){
            cur.next.put(word.charAt(index),new Node());
        }
        cur = cur.next.get(word.charAt(index));
        addRecursive(word,++index,cur);
    }

    /**
     * trie中是否包含单词
     * 非递归写法
     * @param word
     * @return
     */
    public boolean contains(String word){
        Node cur = root;
        for (int i = 0; i < word.length(); i++){
            Character character = word.charAt(i);
            if(null == cur.next.get(character)){
                return false;
            }
            cur = cur.next.get(character);
        }
        return cur.isLeaf;
    }

    /**
     * trie中是否包含单词
     * 递归写法
     * @param word
     * @return
     */
    public boolean containsRecursive(String word){
        return containsRecursive(word,0,root);
    }

    public boolean containsRecursive(String word,int index,Node cur){
        if(index >= word.length()){
            return cur.isLeaf;
        }
        if(null == cur.next.get(word.charAt(index))){
            return false;
        }
        cur = cur.next.get(word.charAt(index));
        return containsRecursive(word,++index,cur);
    }

    public boolean isPrefix(String prefix){
        return isPrefix(prefix,0,root);
    }

    private boolean isPrefix(String prefix, int index, Node cur) {
        if(prefix.length() == index){
            return true;
        }
        if(null == cur.next.get(prefix.charAt(index))){
            return false;
        }
        cur = cur.next.get(prefix.charAt(index));
        return isPrefix(prefix,++index,cur);
    }

    public int getSize(){
        return size;
    }

}
